from collections import OrderedDict
import numpy as np
from typing import Dict, Iterator, Tuple, List, Any, Optional

class Parameter:
    """
    表示可训练参数的包装类
    """
    def __init__(self, data):
        from ml_lib.core import Tensor
        self.data = data if isinstance(data, Tensor) else Tensor(data, requires_grad=True)
        self.requires_grad = True
        
    @property
    def grad(self):
        """获取梯度，代理到内部Tensor的grad"""
        return self.data.grad if hasattr(self.data, 'grad') else None
        
    @grad.setter
    def grad(self, value):
        """设置梯度，代理到内部Tensor的grad"""
        self.data.grad = value
        
    def __repr__(self):
        return f"Parameter({self.data})"
        
    def zero_grad(self):
        """清除梯度"""
        if hasattr(self.data, 'grad'):
            self.data.grad = None

class Module:
    """
    所有神经网络模块的基类
    类似于PyTorch的nn.Module
    """
    def __init__(self):
        self._parameters = OrderedDict()
        self._modules = OrderedDict()
        self.training = True
        
    def __setattr__(self, name, value):
        if isinstance(value, Parameter):
            self._parameters[name] = value
        elif isinstance(value, Module):
            self._modules[name] = value
        super().__setattr__(name, value)
        
    def __call__(self, *args, **kwargs):
        return self.forward(*args, **kwargs)
    
    def parameters(self) -> Iterator[Parameter]:
        """
        返回模块参数的迭代器
        """
        for param in self._parameters.values():
            yield param
            
        for module in self._modules.values():
            for param in module.parameters():
                yield param
    
    def named_parameters(self, prefix: str = '') -> Iterator[Tuple[str, Parameter]]:
        """
        返回模块参数及其名称的迭代器
        
        Args:
            prefix: 参数名前缀
        """
        for name, param in self._parameters.items():
            yield prefix + ('.' if prefix else '') + name, param
            
        for module_name, module in self._modules.items():
            submodule_prefix = prefix + ('.' if prefix else '') + module_name
            for name, param in module.named_parameters(submodule_prefix):
                yield name, param
    
    def train(self, mode: bool = True) -> 'Module':
        """
        将模块设置为训练模式
        
        Args:
            mode: 是否为训练模式
        """
        self.training = mode
        for module in self._modules.values():
            module.train(mode)
        return self
    
    def eval(self) -> 'Module':
        """
        将模块设置为评估模式
        """
        return self.train(False)
    
    def forward(self, *args, **kwargs):
        """
        前向传播，需要在子类中实现
        """
        raise NotImplementedError("子类必须实现forward方法")
    
    def zero_grad(self):
        """
        将所有参数的梯度清零
        """
        for param in self.parameters():
            param.zero_grad()

